from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
import os

os.environ['OPENAI_API_KEY'] = 'sk-ioXCJ3W2IQlpDzHu268aEeF7CcAf4b1f9f8c46A248C74a9d'
os.environ['OPENAI_BASE_URL'] = 'https://api.xiaoai.plus/v1'

if __name__ == '__main__':

    ## 使用提示模板来指导它的回答
    prompt = ChatPromptTemplate.from_messages([
        ("system", "你是一位优秀的数学老师，尤其擅长教授概率相关的知识点"),
        ("user", "{input}")
    ])

    ### 输出解析器，将多余消息过滤，只需要回答字符串
    output_parser = StrOutputParser()

    llm = ChatOpenAI()
    llm.model_name = 'gpt-4'

    ###  组合LLM链
    chain = llm | output_parser

    print(chain.invoke("X,Y独立且均服从均匀分布U(0,1)，求U=X-2Y的密度。"))

